#!/usr/bin/env python3
"""
Transformer音频生成器测试套件
验证音频生成模型的正确性和性能
"""
import torch
import pytest
import sys
import os
import numpy as np
import tempfile

# Add the models directory to the path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'models'))

from models.audio_generator import AudioGenerator

class TestAudioGenerator:
    @pytest.fixture
    def sample_audio_tensor(self):
        """创建测试用的音频张量"""
        # 生成1秒的测试音频，采样率22050
        duration = 1.0  # 1秒
        sample_rate = 22050
        t = torch.linspace(0, duration, int(duration * sample_rate))
        audio = torch.sin(2 * np.pi * 440 * t)  # 440Hz正弦波
        return audio.unsqueeze(0)  # (1, audio_length)
    
    @pytest.fixture
    def audio_generator(self):
        """创建音频生成器实例"""
        return AudioGenerator(
            sample_rate=22050,
            segment_length=1024,
            hop_length=256,
            n_mels=80,
            embedding_dim=512,
            nhead=8,
            num_layers=6,
            max_segments=512
        )

    def test_generator_initialization(self, audio_generator):
        """测试音频生成器正确初始化"""
        assert audio_generator is not None
        assert isinstance(audio_generator, AudioGenerator)
        assert hasattr(audio_generator, 'model')
        assert hasattr(audio_generator, 'mel_transform')
        assert hasattr(audio_generator, 'segment_embedding')
        assert hasattr(audio_generator, 'positional_encoding')
    
    def test_positional_encoding_shape(self, audio_generator):
        """测试位置编码的形状"""
        pos_enc = audio_generator.positional_encoding
        assert pos_enc.shape == (512, 512)  # (max_segments, embedding_dim)
        assert not torch.isnan(pos_enc).any()
        assert not torch.isinf(pos_enc).any()
    
    def test_audio_to_mel_spectrogram(self, audio_generator, sample_audio_tensor):
        """测试音频到梅尔频谱转换"""
        mel_spec = audio_generator._audio_to_mel_spectrogram(sample_audio_tensor)
        
        # 验证形状
        batch_size, time_frames, n_mels = mel_spec.shape
        assert batch_size == 1
        assert n_mels == 80
        
        # 验证时间帧数计算
        audio_length = sample_audio_tensor.shape[1]
        # PyTorch的torch.stft默认使用中心填充，帧数计算公式为：
        # frames = ceil((audio_length) / hop_length)
        expected_frames = (audio_length + audio_generator.hop_length - 1) // audio_generator.hop_length
        # 由于不同的填充策略，允许一定的误差范围
        assert abs(time_frames - expected_frames) <= 2
        
        # 验证数值范围
        assert mel_spec.min() >= 0  # 梅尔频谱应为非负
        assert not torch.isnan(mel_spec).any()
        assert not torch.isinf(mel_spec).any()
    
    def test_mel_spectrogram_to_audio(self, audio_generator):
        """测试梅尔频谱到音频转换"""
        # 创建测试梅尔频谱
        batch_size, time_frames, n_mels = 1, 100, 80
        mel_spec = torch.randn(batch_size, time_frames, n_mels)
        
        audio = audio_generator._mel_spectrogram_to_audio(mel_spec)
        
        # 验证形状
        assert audio.shape[0] == batch_size
        assert not torch.isnan(audio).any()
        assert not torch.isinf(audio).any()
    
    def test_preprocess_audio_tensor(self, audio_generator, sample_audio_tensor):
        """测试张量音频预处理"""
        processed = audio_generator.preprocess_audio(sample_audio_tensor)
        
        assert processed.shape == sample_audio_tensor.shape
        assert torch.max(torch.abs(processed)) <= 1.0  # 标准化检查
        assert not torch.isnan(processed).any()
    
    def test_preprocess_audio_numpy(self, audio_generator, sample_audio_tensor):
        """测试numpy数组音频预处理"""
        audio_numpy = sample_audio_tensor.numpy()
        processed = audio_generator.preprocess_audio(audio_numpy)
        
        assert processed.shape == sample_audio_tensor.shape
        assert torch.max(torch.abs(processed)) <= 1.0
    
    def test_generate_method(self, audio_generator, sample_audio_tensor):
        """测试生成方法"""
        # 使用少量步数进行快速测试
        generated_audio = audio_generator.generate(sample_audio_tensor, num_steps=5, temperature=1.0)
        
        # 验证输出形状和数值
        assert generated_audio.shape[0] == 1  # batch_size
        assert generated_audio.shape[1] > 0  # 音频长度
        assert not torch.isnan(generated_audio).any()
        assert not torch.isinf(generated_audio).any()
    
    def test_generate_continuation(self, audio_generator, sample_audio_tensor):
        """测试音频续写生成"""
        continuation = audio_generator.generate_continuation(
            sample_audio_tensor, 
            continuation_length=1.0,  # 续写1秒
            num_steps=5, 
            temperature=1.0
        )
        
        # 验证续写结果
        assert continuation.shape[0] == 1
        assert continuation.shape[1] > sample_audio_tensor.shape[1]  # 应该比原始音频长
        assert not torch.isnan(continuation).any()
        assert not torch.isinf(continuation).any()
    
    def test_generate_from_texture(self, audio_generator):
        """测试从纹理生成音频"""
        texture_audio = audio_generator.generate_from_texture(duration=1.0)
        
        assert texture_audio.shape[0] == 1
        assert texture_audio.shape[1] > 0
        assert not torch.isnan(texture_audio).any()
        assert not torch.isinf(texture_audio).any()
    
    def test_temperature_effect(self, audio_generator, sample_audio_tensor):
        """测试温度参数对生成的影响"""
        # 低温度：确定性更强
        low_temp_audio = audio_generator.generate(sample_audio_tensor, num_steps=3, temperature=0.1)
        
        # 高温度：随机性更强
        high_temp_audio = audio_generator.generate(sample_audio_tensor, num_steps=3, temperature=2.0)
        
        # 验证两者都有效
        assert low_temp_audio.shape == high_temp_audio.shape
        assert not torch.isnan(low_temp_audio).any()
        assert not torch.isnan(high_temp_audio).any()
    
    def test_batch_processing(self, audio_generator):
        """测试批量音频处理"""
        # 创建批量音频数据
        batch_size = 3
        audio_length = 22050  # 1秒音频
        batch_audio = torch.randn(batch_size, audio_length)
        
        # 批量生成
        generated_batch = audio_generator.generate(batch_audio, num_steps=3)
        
        # 验证批量处理
        assert generated_batch.shape[0] == batch_size
        assert generated_batch.shape[1] > 0
    
    def test_edge_cases(self, audio_generator):
        """测试边界情况"""
        # 测试非常短的音频
        short_audio = torch.randn(1, 512)  # 很短
        short_result = audio_generator.generate(short_audio, num_steps=2)
        assert short_result.shape[0] == 1
        
        # 测试静音音频
        silent_audio = torch.zeros(1, 22050)
        silent_result = audio_generator.generate(silent_audio, num_steps=2)
        assert silent_result.shape[0] == 1
    
    def test_device_compatibility(self, audio_generator, sample_audio_tensor):
        """测试设备兼容性"""
        if torch.cuda.is_available():
            # 测试GPU设备
            audio_generator_gpu = AudioGenerator().cuda()
            audio_gpu = sample_audio_tensor.cuda()
            
            gpu_result = audio_generator_gpu.generate(audio_gpu, num_steps=2)
            assert gpu_result.device.type == 'cuda'
    
    def test_model_forward_pass(self, audio_generator):
        """测试底层模型前向传播"""
        batch_size = 2
        seq_length = 100
        x = torch.randn(seq_length, batch_size, 512)  # embedding_dim=512
        
        output = audio_generator.model(x)
        
        assert output.shape == (seq_length, batch_size, 80)  # n_mels=80
        assert not torch.isnan(output).any()
        assert not torch.isinf(output).any()

if __name__ == "__main__":
    pytest.main([__file__, "-v"])